import json
import random
import copy
import os

# 原始数据文件路径
original_filename = 'assets/user_data/abc.json'

# 加载原始数据
with open(original_filename, 'r') as f:
    original_data = json.load(f)

# 定义更齐全的邻近键映射，只保留 , 和 . 符号
nearby_keys = {
    'q': ['w', 'a'],
    'w': ['q', 'e', 'a', 's'],
    'e': ['w', 'r', 's', 'd'],
    'r': ['e', 't', 'd', 'f'],
    't': ['r', 'y', 'f', 'g'],
    'y': ['t', 'u', 'g', 'h'],
    'u': ['y', 'i', 'h', 'j'],
    'i': ['u', 'o', 'j', 'k'],
    'o': ['i', 'p', 'k', 'l'],
    'p': ['o', 'l'],
    'a': ['q', 'w', 's', 'z'],
    's': ['w', 'e', 'a', 'd', 'z', 'x'],
    'd': ['e', 'r', 's', 'f', 'x', 'c'],
    'f': ['r', 't', 'd', 'g', 'c', 'v'],
    'g': ['t', 'y', 'f', 'h', 'v', 'b'],
    'h': ['y', 'u', 'g', 'j', 'b', 'n'],
    'j': ['u', 'i', 'h', 'k', 'n', 'm'],
    'k': ['i', 'o', 'j', 'l', 'm', ','],
    'l': ['o', 'p', 'k', ',', '.'],
    'z': ['a', 's', 'x'],
    'x': ['s', 'd', 'z', 'c'],
    'c': ['d', 'f', 'x', 'v'],
    'v': ['f', 'g', 'c', 'b'],
    'b': ['g', 'h', 'v', 'n'],
    'n': ['h', 'j', 'b', 'm'],
    'm': ['j', 'k', 'n', ','],
    ',': ['k', 'l', 'm', '.'],
    '.': ['l', ','],
    ' ': ['c', 'v', 'b', 'n', 'm', ',', '.']  # 空格键的邻近键
}

# 定义增强函数
def augment_keystrokes(keystrokes):
    augmented_keystrokes = []
    time_offset = 0.0

    # 调整打字速度（随机缩放因子）
    speed_factor = random.uniform(0.8, 1.2)

    for i, keystroke in enumerate(keystrokes):
        # 第一个按键的 press_time 恒定为 0
        if i == 0:
            press_time = 0.0
        else:
            press_time = keystroke['press_time'] * speed_factor + random.uniform(-0.05, 0.05)

        release_time = keystroke['release_time']
        if release_time is not None:
            release_time = release_time * speed_factor + random.uniform(-0.05, 0.05)
        else:
            # 如果 release_time 为 None，设置一个默认值，比如 press_time + 0.1
            release_time = press_time + 0.1

        # 确保时间不为负且 press_time < release_time
        press_time = max(0, press_time)
        release_time = max(press_time + 0.01, release_time)  # 确保 release_time 大于 press_time

        # 保留四位小数，使用格式化字符串确保严格保留四位
        press_time = format(press_time, '.4f')
        release_time = format(release_time, '.4f')

        # 模拟打字错误（10% 的概率替换按键）
        if random.random() < 0.1 and keystroke['key_content'].lower() in nearby_keys:
            # 替换为邻近键
            key_content = random.choice(nearby_keys[keystroke['key_content'].lower()])
            # 保持原始大小写
            if keystroke['key_content'].isupper():
                key_content = key_content.upper()
        else:
            key_content = keystroke['key_content']

        # 添加增强后的按键，将格式化后的字符串转换回浮点数
        augmented_keystrokes.append({
            'key_content': key_content,
            'press_time': float(press_time),
            'release_time': float(release_time)
        })

        # 模拟打字停顿（20% 的概率）
        if random.random() < 0.2:
            pause_duration = round(random.uniform(0.1, 0.5), 4)
            time_offset += pause_duration

    return augmented_keystrokes

# 生成增强数据
augmented_data = copy.deepcopy(original_data)
original_sequence_count = len(original_data)
new_sequence_count = original_sequence_count * 9  # 固定生成 9 倍的新数据量

for _ in range(new_sequence_count):
    # 随机选择一条原始数据
    original_entry = random.choice(original_data)
    new_entry = {
        'subject': original_entry['subject'],  # 保留原始用户标识
        'keystrokes': augment_keystrokes(original_entry['keystrokes'])
    }
    augmented_data.append(new_entry)

# 保存增强后的数据
base_filename = os.path.basename(original_filename)
output_folder = 'assets/user_data/augmented_user_data'
# 确保输出文件夹存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)
output_filename = os.path.join(output_folder, base_filename)
with open(output_filename, 'w') as f:
    json.dump(augmented_data, f, indent=4)

# 输出统计信息
print(f"原始数据集的序列数: {original_sequence_count}")
print(f"数据增强新添加的序列数: {new_sequence_count}")
print(f"增强后的总序列数: {len(augmented_data)}")
print(f"数据增强完成，已保存到 {output_filename}")    